import os
import time
import shutil
import numpy as np
import tensorflow as tf
import core.utils as utils
from tqdm import tqdm
from core.dataset import Dataset
from core.yolov3 import YOLOV3
from core.config import cfg

class YoloTrain(object):

    def __init__(self):
        #在每一种缩放比例下anchor boxes的数量 3
        self.anchor_per_scale    = cfg.YOLO.ANCHOR_PER_SCALE
        #类名
        self.classes             = utils.read_class_names(cfg.YOLO.CLASSES)
        #类别总数
        self.num_classes         = len(self.classes)
        #初始学习率
        self.learn_rate_init     = cfg.TRAIN.LEARN_RATE_INIT
        #稳定后的学习率
        self.learn_rate_end      = cfg.TRAIN.LEARN_RATE_END
        #第一次epochs的迭代次数
        self.first_stage_epochs  = cfg.TRAIN.FISRT_STAGE_EPOCHS
        #第二次epochs的迭代次数
        self.second_stage_epochs = cfg.TRAIN.SECOND_STAGE_EPOCHS
        #warmup的epochs//TODO:warmup机制的作用？
        self.warmup_periods      = cfg.TRAIN.WARMUP_EPOCHS
        #初始训练权重
        self.initial_weight      = cfg.TRAIN.INITIAL_WEIGHT
        #个数话当前时间
        self.time                = time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime(time.time()))
        #滑动平均衰减率0.9995
        self.moving_ave_decay    = cfg.YOLO.MOVING_AVE_DECAY
        #每种图像缩放比例下最多有150个bbox
        self.max_bbox_per_scale  = 150
        self.train_logdir        = "./data/log/train"
        #训练集
        self.trainset            = Dataset('train')
        #测试集
        self.testset             = Dataset('test')
        #每迭代一次需要喂入的图像个数
        self.steps_per_period    = len(self.trainset)
        #tensorflow会话
        self.sess                = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))

        #定义输入占位
        with tf.name_scope('define_input'):
            self.input_data   = tf.placeholder(dtype=tf.float32, name='input_data')
            self.label_sbbox  = tf.placeholder(dtype=tf.float32, name='label_sbbox')
            self.label_mbbox  = tf.placeholder(dtype=tf.float32, name='label_mbbox')
            self.label_lbbox  = tf.placeholder(dtype=tf.float32, name='label_lbbox')
            self.true_sbboxes = tf.placeholder(dtype=tf.float32, name='sbboxes')
            self.true_mbboxes = tf.placeholder(dtype=tf.float32, name='mbboxes')
            self.true_lbboxes = tf.placeholder(dtype=tf.float32, name='lbboxes')
            self.trainable     = tf.placeholder(dtype=tf.bool, name='training')
        #定义loss值
        with tf.name_scope("define_loss"):
            #YOLO_v3
            self.model = YOLOV3(self.input_data, self.trainable)
            #网络变量
            self.net_var = tf.global_variables()
            #计算不同种类的loss值
            self.giou_loss, self.conf_loss, self.prob_loss = self.model.compute_loss(
                                                    self.label_sbbox,  self.label_mbbox,  self.label_lbbox,
                                                    self.true_sbboxes, self.true_mbboxes, self.true_lbboxes)
            #总loss
            self.loss = self.giou_loss + self.conf_loss + self.prob_loss
        #学习率
        with tf.name_scope('learn_rate'):
            #定义global_step
            self.global_step = tf.Variable(1.0, dtype=tf.float64, trainable=False, name='global_step')
            #warmup迭代总次数=warmup的epoch*每迭代一次需要喂入的图像个数
            warmup_steps = tf.constant(self.warmup_periods * self.steps_per_period,
                                        dtype=tf.float64, name='warmup_steps')
            #训练迭代总次数=(第一次epoch+第二次epoch)*每迭代一次需要喂入的图像个数
            train_steps = tf.constant( (self.first_stage_epochs + self.second_stage_epochs)* self.steps_per_period,
                                        dtype=tf.float64, name='train_steps')
            #当pred为true时执行true_fn,当pred为false时执行false_fn
            self.learn_rate = tf.cond(
                pred=self.global_step < warmup_steps,
                true_fn=lambda: self.global_step / warmup_steps * self.learn_rate_init,
                false_fn=lambda: self.learn_rate_end + 0.5 * (self.learn_rate_init - self.learn_rate_end) *
                                    (1 + tf.cos(
                                        (self.global_step - warmup_steps) / (train_steps - warmup_steps) * np.pi))
            )
            #更新global_step
            global_step_update = tf.assign_add(self.global_step, 1.0)
        #将滑动平均模型应用于所有训练的变量，可以使得模型在测试数据集上更加健壮
        with tf.name_scope("define_weight_decay"):
            moving_ave = tf.train.ExponentialMovingAverage(self.moving_ave_decay).apply(tf.trainable_variables())
        with tf.name_scope("define_first_stage_train"):
            #训练变量集合
            self.first_stage_trainable_var_list = []
            for var in tf.trainable_variables():
                var_name = var.op.name
                var_name_mess = str(var_name).split('/')
                if var_name_mess[0] in ['conv_sbbox', 'conv_mbbox', 'conv_lbbox']:
                    self.first_stage_trainable_var_list.append(var)
            #adam优化器
            first_stage_optimizer = tf.train.AdamOptimizer(self.learn_rate).minimize(self.loss,
                                                      var_list=self.first_stage_trainable_var_list)

            with tf.control_dependencies(tf.get_collection(tf.GraphKeys.UPDATE_OPS)):
                with tf.control_dependencies([first_stage_optimizer, global_step_update]):
                    with tf.control_dependencies([moving_ave]):
                        self.train_op_with_frozen_variables = tf.no_op()

        with tf.name_scope("define_second_stage_train"):
            second_stage_trainable_var_list = tf.trainable_variables()
            second_stage_optimizer = tf.train.AdamOptimizer(self.learn_rate).minimize(self.loss,
                                                      var_list=second_stage_trainable_var_list)

            with tf.control_dependencies(tf.get_collection(tf.GraphKeys.UPDATE_OPS)):
                with tf.control_dependencies([second_stage_optimizer, global_step_update]):
                    with tf.control_dependencies([moving_ave]):
                        self.train_op_with_all_variables = tf.no_op()

        with tf.name_scope('loader_and_saver'):
            #保存所有变量
            self.loader = tf.train.Saver(self.net_var)
            #max_to_keep:指定要保留checkpoint的最大数量
            self.saver  = tf.train.Saver(tf.global_variables(), max_to_keep=10)

        #日志信息
        with tf.name_scope('summary'):
            tf.summary.scalar("learn_rate", self.learn_rate)
            tf.summary.scalar("giou_loss",  self.giou_loss)
            tf.summary.scalar("conf_loss",  self.conf_loss)
            tf.summary.scalar("prob_loss",  self.prob_loss)
            tf.summary.scalar("total_loss", self.loss)

            logdir = "./data/log/"
            if os.path.exists(logdir): shutil.rmtree(logdir)
            os.mkdir(logdir)
            self.write_op = tf.summary.merge_all()
            self.summary_writer  = tf.summary.FileWriter(logdir, graph=self.sess.graph)

    def train(self):
        self.sess.run(tf.global_variables_initializer())
        try:
            print('=> Restoring weights from: %s ... ' % self.initial_weight)
            #加载已保存的模型，模型的初始化参数为self.intial_weight
            self.loader.restore(self.sess, self.initial_weight)
        except:
            print('=> %s does not exist !!!' % self.initial_weight)
            print('=> Now it starts to train YOLOV3 from scratch ...')
            self.first_stage_epochs = 0

        for epoch in range(1, 1+self.first_stage_epochs+self.second_stage_epochs):
            if epoch <= self.first_stage_epochs:
                train_op = self.train_op_with_frozen_variables
            else:
                train_op = self.train_op_with_all_variables
            #用于显示trainset加载的进度条
            pbar = tqdm(self.trainset)
            #一个epoch后训练与测试的误差
            train_epoch_loss, test_epoch_loss = [], []
            #循环一次，就会调用一次dataset.py的__next__()方法,返回值以tuple类型赋予train_data
            for train_data in pbar:
                _, summary, train_step_loss, global_step_val = self.sess.run(
                    [train_op, self.write_op, self.loss, self.global_step],feed_dict={
                                                self.input_data:   train_data[0],
                                                self.label_sbbox:  train_data[1],
                                                self.label_mbbox:  train_data[2],
                                                self.label_lbbox:  train_data[3],
                                                self.true_sbboxes: train_data[4],
                                                self.true_mbboxes: train_data[5],
                                                self.true_lbboxes: train_data[6],
                                                self.trainable:    True,
                })
                #训练集误差list
                train_epoch_loss.append(train_step_loss)
                self.summary_writer.add_summary(summary, global_step_val)
                #进度条说明
                pbar.set_description("train loss: %.2f" %train_step_loss)

            for test_data in self.testset:
                test_step_loss = self.sess.run( self.loss, feed_dict={
                                                self.input_data:   test_data[0],
                                                self.label_sbbox:  test_data[1],
                                                self.label_mbbox:  test_data[2],
                                                self.label_lbbox:  test_data[3],
                                                self.true_sbboxes: test_data[4],
                                                self.true_mbboxes: test_data[5],
                                                self.true_lbboxes: test_data[6],
                                                self.trainable:    False,
                })
                #测试集误差list
                test_epoch_loss.append(test_step_loss)
            #训练集、测试集平均误差
            train_epoch_loss, test_epoch_loss = np.mean(train_epoch_loss), np.mean(test_epoch_loss)
            #checkpoint文件命名
            ckpt_file = "./checkpoint/yolov3_test_loss=%.4f.ckpt" % test_epoch_loss
            #获取当前时间
            log_time = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time()))

            print("=> Epoch: %2d Time: %s Train loss: %.2f Test loss: %.2f Saving %s ..."
                            %(epoch, log_time, train_epoch_loss, test_epoch_loss, ckpt_file))
            #保存训练模型
            self.saver.save(self.sess, ckpt_file, global_step=epoch)



if __name__ == '__main__':
    YoloTrain().train()




